using MultiAgentsClient.Service.WorkflowEngine.BaseNodes;
using MultiAgentsClient.Service.WorkflowEngine.BuiltInNodes;

namespace MultiAgentsClient.Service.WorkflowEngine;

public class BehaviorTreeSession
{
    private readonly BehaviorTreeContext Context;
    public BehaviorTreeConfig behaviorTreeConfig;

    public BehaviorTreeSession(int chatRoomIdIn, BehaviorTreeConfig configIn)
    {
        Context = new BehaviorTreeContext()
        {
            ChatRoomIdForSession = chatRoomIdIn,
            BehaviorTreeName = configIn.name,
        };

        behaviorTreeConfig = BehaviorTreeConfig.MakeNewCopyFromJson(configIn.workflowId, configIn.ownerId, configIn.name, configIn.configJson);

        behaviorTreeConfig.TriggerNode.InitNodeForRuntime(Context);
        foreach (BaseNode opNode in behaviorTreeConfig.OperationNodes)
        {
            opNode.InitNodeForRuntime(Context);
        }
    }

    public async Task RunAsync()
    {
        BaseNode? currentNode = behaviorTreeConfig.TriggerNode;
        await HandleNode(currentNode);
    }

    private async Task<NodeOutcome> HandleNode(BaseNode currentNode)
    {
        if (currentNode is ForEachNode foreachNode)
        {
            return await HandleForEachNode(foreachNode);
        }
        
        if (currentNode is AndNode || currentNode is OrNode)
        {
            return await HandleAndOrNode(currentNode);
        }

        var curNodeOutcome = await currentNode.ExecuteNodeAsync();

        if (curNodeOutcome.NextNodeIndex < 0)
        {
            return curNodeOutcome;
        }

        //Get next node to execute
        var nextNode = GetNextNode(currentNode, curNodeOutcome);

        // Recursively handle the next node
        if (nextNode != null)
        {
            return await HandleNode(nextNode);
        }

        // Return the current node outcome if there's no next node
        return curNodeOutcome;

    }

    private async Task<NodeOutcome> HandleAndOrNode(BaseNode currentNode)
    {
        if (currentNode is AndNode sequenceNode)
        {
            var nextNodes = GetSubNodesSorted(sequenceNode);
            foreach (var nextNode in nextNodes)
            {
                NodeOutcome subNodeOutcome = await HandleNode(nextNode);
                if (subNodeOutcome.Status != NodeOutcomeStatus.Success) 
                {
                    //any failure would result in termination
                    return subNodeOutcome;
                }
            }

            return new NodeOutcome(NodeOutcomeStatus.Success, -1);
        }
        else if (currentNode is OrNode selectorNode)
        {
            var nextNodes = GetSubNodesSorted(selectorNode);
            foreach (var nextNode in nextNodes)
            {
                NodeOutcome subNodeOutcome = await HandleNode(nextNode);
                if (subNodeOutcome.Status == NodeOutcomeStatus.Success)
                {
                    return subNodeOutcome;
                }
            }

            return new NodeOutcome(NodeOutcomeStatus.Failure, -1);
        }

        throw new Exception("Should not happen");
    }


    private async Task<NodeOutcome> HandleForEachNode(ForEachNode foreachNode)
    {
        NodeOutcome foreachOutcome;

        while (true)
        {
            foreachOutcome = await foreachNode.ExecuteNodeAsync();

            if (foreachOutcome.NextNodeIndex == -1) // Iteration completed
            {
                break;
            }

            // Get the next node connected to the specific output port
            var nextNode = GetNextNode(foreachNode, foreachOutcome);

            if (nextNode != null)
            {
                NodeOutcome execOutcome = await HandleNode(nextNode);
                // This allows continuing the chain correctly
                while (execOutcome.NextNodeIndex != -1)
                {
                    nextNode = GetNextNode(nextNode, execOutcome);
                    if (nextNode == null) break;
                    execOutcome = await nextNode.ExecuteNodeAsync();
                }
            }
        }

        return foreachOutcome;

        //// After ForEachNode is done, get the next node (if any) for sequential processing
        //var nextSequentialNode = GetNextNode(foreachNode, 0); // Assuming next sequential node is connected to port 0
        //if (nextSequentialNode != null)
        //{
        //    HandleNodeChain(nextSequentialNode);
        //}
    }

    private List<BaseNode> GetSubNodesSorted(BaseNode baseNode)
    {
        List<BaseNode> connectedNodes = new List<BaseNode>();

        foreach (var connection in behaviorTreeConfig.connections)
        {
            if (connection.FromNodeId == baseNode.Id)
            {
                BaseNode? toNode = behaviorTreeConfig.OperationNodes.FirstOrDefault(n => n.Id == connection.ToNodeId);
                if (toNode != null)
                {
                    connectedNodes.Add(toNode);
                }
            }
        }

        return connectedNodes.OrderBy(n => n.nodePosY).ToList();
    }

    private BaseNode? GetNextNode(BaseNode curNode, NodeOutcome curNodeOutput)
    {
        WorkflowConnection hookedConnection = null;
        //find connection
        foreach (var connection in behaviorTreeConfig.connections)
        {
            if (connection.FromNodeId == curNode.Id && connection.FromNodeOutputPortIndex == curNodeOutput.NextNodeIndex)
            {
                hookedConnection = connection;
            }
        }

        if (hookedConnection == null)
        {
            return null;
        }

        //find node
        foreach (var operationNode in behaviorTreeConfig.OperationNodes)
        {
            if (hookedConnection.ToNodeId == operationNode.Id)//TODO:  && hookedConnection.ToNodeInputPortIndex == portIndex)
            {
                return operationNode;
            }
        }
        return null;
    }
}
